#include "nes.h"
#include "nes_mapper.h"
#include <string.h>

typedef struct {
    uint8_t shift_register;
    uint8_t shift_count;
    uint8_t control;
    uint8_t chr_bank0;
    uint8_t chr_bank1;
    uint8_t prg_bank;
} mmc1_state_t;

static mmc1_state_t mmc1;

static void mmc1_update_banks(nes_t* nes) {
    uint8_t prg_mode = (mmc1.control >> 2) & 0x03;
    uint8_t prg_bank = mmc1.prg_bank & 0x0F;

    switch (prg_mode) {
        case 0:
        case 1:
            // 32 KB mode, ignore LSB
            prg_bank &= 0x0E;
            nes_load_prgrom_16k(nes, 0, prg_bank);
            nes_load_prgrom_16k(nes, 1, prg_bank + 1);
            break;
        case 2:
            // Fixed first bank at $8000, switch $C000
            nes_load_prgrom_16k(nes, 0, 0);
            nes_load_prgrom_16k(nes, 1, prg_bank);
            break;
        case 3:
            // Switch $8000, fixed last bank at $C000
            nes_load_prgrom_16k(nes, 0, prg_bank);
            nes_load_prgrom_16k(nes, 1, nes->nes_rom.prg_rom_size - 1);
            break;
    }

    // CHR bank switching (assume 8KB CHR)
    nes_load_chrrom_8k(nes, 0, mmc1.chr_bank0);
}

static void mmc1_write(nes_t* nes, uint16_t addr, uint8_t data) {
    if (addr < 0x8000) return;

    // Reset shift register if bit 7 set
    if (data & 0x80) {
        mmc1.shift_register = 0x10;
        mmc1.shift_count = 0;
        mmc1.control |= 0x0C;  // Default to 16KB PRG, last bank fixed
        mmc1_update_banks(nes);
        return;
    }

    // Load into shift register
    mmc1.shift_register = (mmc1.shift_register >> 1) | ((data & 1) << 4);
    mmc1.shift_count++;

    if (mmc1.shift_count == 5) {
        uint16_t bank = (addr >> 13) & 0x03;
        switch (bank) {
            case 0: mmc1.control = mmc1.shift_register; break;
            case 1: mmc1.chr_bank0 = mmc1.shift_register; break;
            case 2: mmc1.chr_bank1 = mmc1.shift_register; break;
            case 3: mmc1.prg_bank = mmc1.shift_register; break;
        }

        mmc1_update_banks(nes);
        mmc1.shift_register = 0x10;
        mmc1.shift_count = 0;
    }
}

static void mmc1_init(nes_t* nes) {
    memset(&mmc1, 0, sizeof(mmc1));
    mmc1.shift_register = 0x10;
    mmc1.control = 0x0C;  // default mode: PRG-ROM switch at $8000, fixed $C000
    mmc1_update_banks(nes);
}

int nes_mapper1_init(nes_t* nes) {
    nes->nes_mapper.mapper_init = mmc1_init;
    nes->nes_mapper.mapper_write = mmc1_write;
    nes->nes_mapper.mapper_sram = NULL;
    nes->nes_mapper.mapper_apu = NULL;
    nes->nes_mapper.mapper_read_apu = NULL;
    nes->nes_mapper.mapper_vsync = NULL;
    nes->nes_mapper.mapper_hsync = NULL;
    nes->nes_mapper.mapper_ppu = NULL;
    nes->nes_mapper.mapper_render_screen = NULL;
    return 0;
}
